package com.gitee.wsl.mathematics.function.randomForest

import kotlin.math.abs
import kotlin.math.ln
import kotlin.random.Random

/**
 * @author lidapeng
 * @description
 * @date 3:12 下午 2020/2/17
 */
class Tree {
    //决策树
    var dataTable: DataTable? = null
        private set
    private var table: Map<String, List<Int>>? = null //总样本
    var rootNode: Node? = null //根节点
    private var endList: List<Int>? = null //最终结果分类
    private val lastNodes: MutableList<Node> = ArrayList<Node>() //最后一层节点集合
    private val random: Random = Random.Default
    private val trustPunishment: Float //信任惩罚

    private class Gain {
        var gain: Float = 0f //信息增益
        var gainRatio: Float = 0f //信息增益率
    }

    constructor(trustPunishment: Float) {
        this.trustPunishment = trustPunishment
    }

    constructor(dataTable: DataTable?, trustPunishment: Float) {
        if (dataTable?.key != null) {
            this.trustPunishment = trustPunishment
            this.dataTable = dataTable
        } else {
            throw Exception("dataTable is empty")
        }
    }

    private fun log2(p: Float): Float {
        return ln(p) / ln(2.0f)
    }

    private fun getEnt(list: List<Int>): Float {
        //记录了每个类别有几个
        val myType: MutableMap<Int, Int> = HashMap<Int, Int>()
        for (index in list) {
            val type = endList!![index] //最终结果的类别
            if (myType.containsKey(type)) {
                myType[type] = myType[type]!! + 1
            } else {
                myType[type] = 1
            }
        }
        var ent = 0f
        //求信息熵
        for ((_, value) in myType) {
            val g = value.toFloat() / list.size.toFloat() //每个类别的概率
            ent += g * log2(g)
        }
        return -ent
    }

    private fun getGain(ent: Float, dNub: Float, gain: Float): Float {
        return gain + ent * dNub
    }

    private fun createNode(node: Node): List<Node>? {
        val attributes = node.attribute
        val fatherList = node.fatherList
        if (attributes!!.isNotEmpty()) {
            val mapAll: MutableMap<String, MutableMap<Int, MutableList<Int>>> = HashMap() //主键：可用属性 次主键该属性的属性值，集合该值对应的id
            val fatherEnt = getEnt(fatherList!!)
            val fatherNub = fatherList.size //总样本数
            //该属性每个离散数据分类的集合
            for (i in fatherList.indices) {
                val index = fatherList[i] //编号
                for (attr in attributes) {
                    if (!mapAll.containsKey(attr)) {
                        mapAll[attr] = HashMap()
                    }
                    val map =
                        mapAll[attr]!!
                    val attrValue = table!![attr]!![index] //获取当前属性值
                    if (!map.containsKey(attrValue)) {
                        map[attrValue] = ArrayList<Int>()
                    }
                    val list = map[attrValue]!!
                    list.add(index)
                }
            }
            val nodeMap: MutableMap<String?, List<Node>> = HashMap()
            var i = 0
            var sigmaG = 0f
            val gainMap: MutableMap<String, Gain> = HashMap<String, Gain>()
            for ((name, map) in mapAll) {
                //求信息增益
                var gain = 0f //信息增益
                var IV = 0f //增益率
                val nodeList: MutableList<Node> = ArrayList<Node>()
                nodeMap[name] = nodeList
                for ((key, list) in map) { //遍历当前属性下的所有属性值的集合
                    val nowAttribute = removeAttribute(
                        attributes, name
                    )
                    val sonNode = Node()
                    nodeList.add(sonNode)
                    sonNode.attribute = nowAttribute
                    sonNode.fatherList = list
                    sonNode.typeId = key //该属性值
                    val myNub = list.size //该属性值下数据的数量
                    val ent = getEnt(list) //该属性值 的信息熵
                    val dNub = myNub.toFloat() / fatherNub.toFloat() //该属性值在 父级样本中出现的概率
                    IV += dNub * log2(dNub)
                    gain = getGain(ent, dNub, gain)
                }
                val gain1 = Gain()
                gainMap[name] = gain1
                gain1.gain = fatherEnt - gain //信息增益
                if (IV != 0f) {
                    gain1.gainRatio = gain1.gain / -IV //增益率
                } else {
                    gain1.gainRatio = 1000000f
                }
                sigmaG += gain1.gain
                i++
            }
            val avgGain = sigmaG / i
            var gainRatio = -2f //最大增益率
            var key: String? = null //可选属性
            //System.out.println("平均信息增益==============================" + avgGain);
            for ((key1, gain) in gainMap) {
                //                System.out.println("主键:" + entry.getKey() + ",平均信息增益:" + avgGain + ",可用属性数量:" + gainMap.size()
//                        + "该属性信息增益:" + gain.gain + ",该属性增益率：" + gain.gainRatio + ",当前最高增益率:" + gainRatio);
                if (gainMap.size == 1 || ((gain.gain >= avgGain || abs((gain.gain - avgGain)) < 0.000001) && (gain.gainRatio >= gainRatio || gainRatio == -2f))) {
                    gainRatio = gain.gainRatio
                    key = key1
                }
            }
            node.key = key
            val nodeList =
                nodeMap[key]!!
            for (j in nodeList.indices) { //儿子绑定父亲关系
                nodeList[j].fatherNode = node
            }
            for (j in nodeList.indices) {
                val node1 = nodeList[j]
                node1.nodeList = createNode(node1)
            }
            return nodeList
        } else {
            //判断类别
            node.isEnd = true //叶子节点
            node.type = getType(fatherList!!)
            lastNodes.add(node) //将全部最后一层节点集合
            return null
        }
    }

    private fun getType(list: List<Int>): Int {
        val myType: MutableMap<Int, Int> = HashMap<Int, Int>()
        for (index in list) {
            val type = endList!![index] //最终结果的类别
            if (myType.containsKey(type)) {
                myType[type] = myType[type]!! + 1
            } else {
                myType[type] = 1
            }
        }
        var type = 0
        var nub = 0
        for ((key, nowNub) in myType) {
            if (nowNub > nub) {
                type = key
                nub = nowNub
            }
        }
        return type
    }

    private fun removeAttribute(attributes: Set<String>, name: String): Set<String> {
        val attriBute: MutableSet<String> = HashSet<String>()
        for (myName in attributes) {
            if (myName != name) {
                attriBute.add(myName)
            }
        }
        return attriBute
    }

    @Throws(Exception::class)
    private fun getTypeId(ob: Any, name: String): Int {
        //val body: Class<*> = ob.javaClass
        //val methodName = "get" + name.substring(0, 1).uppercase(Locale.getDefault()) + name.substring(1)
        //val method: reflect.Method = body.getMethod(methodName)
        //return method.invoke(ob).toString().toInt()
        TODO()
    }

    @Throws(Exception::class)
    fun judge(ob: Any): TreeWithTrust { //进行类别判断
        if (rootNode != null) {
            val treeWithTrust = TreeWithTrust()
            treeWithTrust.trust = 1.0f
            goTree(ob, rootNode!!, treeWithTrust, 0)
            return treeWithTrust
        } else {
            throw Exception("rootNode is null")
        }
    }

    private fun punishment(treeWithTrust: TreeWithTrust) { //信任惩罚
        //System.out.println("惩罚");
        var trust = treeWithTrust.trust //获取当前信任值
        trust *= trustPunishment
        treeWithTrust.trust = trust
    }

    @Throws(Exception::class)
    private fun goTree(ob: Any, node: Node, treeWithTrust: TreeWithTrust, times: Int) { //从树顶向下攀爬
        var node = node
        var times = times
        if (!node.isEnd) {
            val myType = getTypeId(ob, node.key!!) //当前类别的ID
            if (myType == 0) { //做信任惩罚
                punishment(treeWithTrust)
            }
            val nodeList = node.nodeList
            var isOk = false
            for (testNode in nodeList!!) {
                if (testNode.typeId == myType) {
                    isOk = true
                    node = testNode
                    break
                }
            }
            if (!isOk) { //当前类别缺失，未知的属性值
                punishment(treeWithTrust)
                val index: Int = random.nextInt(nodeList.size)
                node = nodeList[index]
            }
            times++
            goTree(ob, node, treeWithTrust, times)
        } else {
            //当以0作为结束的时候要做严厉的信任惩罚
            if (node.typeId == 0) {
                val nub = rootNode!!.attribute!!.size - times
                //System.out.println("惩罚次数" + nub);
                for (i in 0 until nub) {
                    punishment(treeWithTrust)
                }
            }
            treeWithTrust.type = node.type
        }
    }

    @Throws(Exception::class)
    fun study() {
        if (dataTable != null && dataTable!!.length > 0) {
            rootNode = Node()
            table = dataTable!!.table
            endList = dataTable!!.table[dataTable!!.key]
            val set  = dataTable!!.keyType
            set.remove(dataTable!!.key)
            rootNode!!.attribute = set //当前可用属性
            val list: MutableList<Int> = ArrayList<Int>()
            for (i in endList!!.indices) {
                list.add(i)
            }
            rootNode!!.fatherList = list //当前父级样本
            rootNode!!.nodeList = createNode(rootNode!!)
            //进行后剪枝
            for (lastNode in lastNodes) {
                prune(lastNode.fatherNode)
            }
            lastNodes.clear()
        } else {
            throw Exception("dataTable is null")
        }
    }

    private fun prune(node: Node?) { //执行剪枝
        if (node != null && !node.isEnd) {
            val listNode = node.nodeList //子节点
            if (isPrune(node, listNode!!)) { //剪枝
                deduction(node)
                prune(node.fatherNode)
            }
        }
    }

    private fun deduction(node: Node) {
        node.isEnd = true
        node.nodeList = null
        node.type = getType(node.fatherList!!)
    }

    private fun isPrune(father: Node, sonNodes: List<Node>): Boolean {
        var isRemove = false
        val typeList: MutableList<Int> = ArrayList<Int>()
        for (i in sonNodes.indices) {
            val node = sonNodes[i]
            val list = node.fatherList
            typeList.add(getType(list!!))
        }
        val fatherType = getType(father.fatherList!!)
        val nub = getRightPoint(father.fatherList!!, fatherType)
        //父级该样本正确率
        val rightFather = nub.toFloat() / father.fatherList!!.size.toFloat()
        var rightNub = 0
        var rightAllNub = 0
        for (i in sonNodes.indices) {
            val node = sonNodes[i]
            val list = node.fatherList
            val right = getRightPoint(list!!, typeList[i])
            rightNub += right
            rightAllNub += list.size
        }
        val rightPoint = rightNub.toFloat() / rightAllNub.toFloat() //子节点正确率
        if (rightPoint <= rightFather) {
            isRemove = true
        }
        return isRemove
    }

    private fun getRightPoint(types: List<Int>, type: Int): Int {
        var nub = 0
        for (index in types) {
            val end = endList!![index]
            if (end == type) {
                nub++
            }
        }
        return nub
    }
}
